#include "krtpch.h"

#include "../../include/Core/RayTracer.h"
#include "../../include/Utility/KRTMath.h"

#include <ctime>

namespace krt {

	RayTracer::RayTracer(Ref<Scene> scene, int samplerNumbers, int depth)
		: m_Scene(scene), m_SamplerNumbers(samplerNumbers), m_Depth(depth) {

		int coreNum;
#if defined KRT_PLATFORM_WINDOWS
		SYSTEM_INFO systemInfo;
		GetSystemInfo(&systemInfo);
		coreNum = (int)(systemInfo.dwNumberOfProcessors);
#else 
		coreNum = 2;
#endif // KRT_PLATFORM_WINDOWS

		m_State = MakeRef<TraceState>();

		m_State->threadNum = coreNum - 2;

		m_ThreadPool = MakeRef<ThreadPool>(m_State->threadNum);

	}

	RayTracer::~RayTracer() {
		// state should disconstruct AFTER thread pool!!!
	}

	void RayTracer::Start(ImageBuffer& imgBuffer, bool isMulti_Thread, int unit_height) {
		isTracing = true;
		int width = m_Scene->GetSize().x, height = m_Scene->GetSize().y;
		if (isMulti_Thread) {
			int splitNum = height / unit_height;
			m_State->taskNum = height % unit_height ? splitNum + 1 : splitNum;
			int tmpHeight = height;
			for (int i = 0; i < m_State->taskNum; i++) {
				if (tmpHeight < unit_height) {
					m_State->progress.emplace_back(0, width * tmpHeight * m_SamplerNumbers);
				}
				else {
					m_State->progress.emplace_back(0, width * unit_height * m_SamplerNumbers);
				}
				tmpHeight -= unit_height;
			}
			int h_start = 0;
			for (int i = 0; i < m_State->taskNum; i++) {
				m_ThreadPool->Submit(traceTask, MakeRef<TaskParam>(
					TaskParam{ *m_Scene, imgBuffer, m_State->progress[i], width, height,
					h_start, h_start += unit_height, m_SamplerNumbers, m_Depth,isTracing }));
			}
		}
		else {
			m_State->taskNum = 1;
			m_State->progress.push_back({ 0, width * height * m_SamplerNumbers });

			m_ThreadPool->Submit(traceTask, MakeRef<TaskParam>(TaskParam{
				*m_Scene, imgBuffer, m_State->progress[0], width, height, 0, height, m_SamplerNumbers, m_Depth,isTracing }));
		}
	}

	void RayTracer::traceTask(Ref<TaskParam> param) {
		auto uv = glm::vec2(0.0f);
		auto fragColor = glm::vec3(0.0f);
		if (!param->isTracing)return;
		for (int h = param->start_h; h < param->height && h < param->end_h; h++) {
			for (int w = 0; w < param->width; w++) {
				for (int s = 0; s < param->sampler; s++) {
					if (!param->isTracing)return;
					// Get uv
					uv = {
						float(w + KRTMath::Drand48()) / float(param->width),
						float(h + KRTMath::Drand48()) / float(param->height)
					};
					// Trace ray
					fragColor += fragTracer(param->scene, uv, param->depth);
					param->progress.current++;
				}
				fragColor /= (float)param->sampler;
				//  gamma 2.2
				fragColor.x = glm::pow(fragColor.x, 0.455f);
				fragColor.y = glm::pow(fragColor.y, 0.455f);
				fragColor.z = glm::pow(fragColor.z, 0.455f);
				// Write frag color
				param->imgBuffer.WritePixel(glm::ivec2(w, h), fragColor);
				fragColor = glm::vec3(0.0f);
			}
		}
	}

	glm::vec3 RayTracer::fragTracer(const Scene& scene, const glm::vec2& uv, int depth) {
		auto re_attenuation = glm::vec3(1.0f);
		Intersection rec;
		Ray scatttered = scene.GetCamera()->GetRay(uv);
		while (depth--) {
			if (scene.GetObjects()->Intersect(scatttered, rec)) {
				if (!rec.material->Scatter(scatttered, rec, re_attenuation, scatttered)) {
					return re_attenuation * rec.material->Emitted(rec.uv);
				}
			}
			else {
				return re_attenuation * scene.GetSky(glm::normalize(scatttered.direct));
			}
		}
		return { 0.0f,0.0f,0.0f };
	}

}
